package geppetto.phraseHMM;

import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.corpus.LanguageVocab;
import gnu.trove.TIntObjectHashMap;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.LinkedList;



//Contains all possible phrases extracted form a corpus.
public class WordTrie {

	PhraseNode _root;
	public int _phraseCounter;
	TIntObjectHashMap<PhraseNode> _phrases;
	int _maxPhraseLen;
	int[] _phrasesByLen;
	
	public WordTrie(int maxPhraseLen) {
		System.out.println("Building word Trie for phrase Len of " + maxPhraseLen);	
		_root = new PhraseNode(-1,-1,null);
		_phrases = new TIntObjectHashMap<PhraseNode>();
		_phraseCounter = 0;
		this._maxPhraseLen = maxPhraseLen;
		_phrasesByLen = new int[maxPhraseLen];
		
	}
	
	public void compactTrie(){
		System.out.println("Compacting trie");
		_phrases.compact();
		//_root.compact(); better to do it while adding or memory will run out first
	}
	
	public PhraseNode insertPhrase(int lastWordId, PhraseNode parent, int count ,int phraseId){
		PhraseNode node = parent.getChild(lastWordId);
		if(node == null){
			_phraseCounter=Math.max(phraseId, _phraseCounter);
			node = parent.addChild(lastWordId, phraseId);
			_phrases.put(phraseId, node);
			_phrasesByLen[node._level-1]++;
		}
		node.count += count;

		return node;
	}
	
	// Inserts a new phrase in case the phrase does not exist or adds the count to thar phrase
	// It he prefix phrase does not exist then the phrase is not inserted
	public PhraseNode insertPhrase(int[] phrase, int count, int phraseId) {	
		//System.out.println(intArrayToString(phrase));
		PhraseNode parent = getPhrase(phrase, phrase.length-1);
		if(parent == null){
			System.out.println("Got null parent ");
			return null; //No prefix		
		}
		
		//Has prefix
		int lastWordId = phrase[phrase.length-1];
		return insertPhrase(lastWordId, parent, count, phraseId);
	}
		
	public PhraseNode insertPhrase(int[] phrase, int count) {	
		return insertPhrase(phrase, count, _phraseCounter+1);
	}
	
	public PhraseNode insertPhrase(int[] phrase) {
		return insertPhrase(phrase, 1);
	}
	
	public PhraseNode insertPhraseRecursive(int[] phrase){
		return insertPhraseRecursive(phrase, 0);
	}
	
	public PhraseNode insertPhraseRecursive(int[] phrase, int count){
		PhraseNode node = _root;
		int depth = phrase.length;
		for(int i = 0; i < depth; i++){		
			if(!node.containsChild(phrase[i])){
				insertPhrase(phrase[i], node, 0, _phraseCounter + 1);
			}
			node = node.getChild(phrase[i]);
		}
		node.setCount(node.getCount()+count);
		return node;

/*		PhraseNode parent = getPhrase(phrase, phrase.length-1);
		if(parent == null){
			int[] smallerPhrase = Arrays.copyOf(phrase, phrase.length-1);
			insertPhraseRecursive(smallerPhrase, 0);
		}
		return insertPhrase(phrase, count);*/
	}
	
	
	public PhraseNode getPhrase(int[] phrase) {
		return getPhrase(phrase, phrase.length);
	}

	
	public PhraseNode getPhrase(int[] phrase, int depth) {
		PhraseNode node = _root;
		for(int i = 0; i < depth; i++){
			node = node.getChild(phrase[i]);
			if(node == null){
				return null; //Does not have prefix
			}
		}
		return node;
	}
	
	public double getPhraseFrequency(int[] phrase) {
		return ((double)getPhrase(phrase).getCount()) / _phraseCounter;
	}
	
	public int getPhraseId(int[] phrase) {
		PhraseNode p = getPhrase(phrase);
		//System.out.println((p == null) + "["+intArrayToString(phrase)+"]");
		if(p == null) {
			return -1;
		}
		return p._phraseId;
	}

	
	public int[] getPhraseIndexesById(int phraseId) {
		PhraseNode node = _phrases.get(phraseId);
		if (node == null) {
			System.out.println("Phrase Tree " + this);
			System.out.println("WordTrie: Failed to find phrase with Id " + phraseId);
			System.exit(-1);
		}
		int[] phrase = new int[node._level];
		for (int i = node._level; i > 0; i--) {
			phrase[node._level - 1] = node._wordId;
			node = node._parent;
		}
		return phrase;
	}
	
	public int[] getPhraseIndexes(PhraseNode node) {
		return getPhraseIndexesById(node._phraseId);
	}
	
	public String toString() {
		return _root.toString() + "\n";
	
	}

	
	public  void addPhrasesAllSentences(ArrayList<int[]> sentences){
		for (int i = 0; i < sentences.size(); i++) {
			int[] sentence = sentences.get(i);
			for (int index = 0; index < sentence.length; index++) {
				int phraseStartIndex = index;
				for (int duration = 1; duration <= _maxPhraseLen; duration++) {
					int phraseEndIndex = phraseStartIndex
							+ duration;
					if (phraseEndIndex > sentence.length)
						break;
					int[] phrase = new int[duration];
					for (int s = 0; s < duration; s++) {
						phrase[s] = sentence[index + s];
					}
					insertPhrase(phrase);
				}
			}
		}
	}

	//Gets all phrases of a given lenght
	public void addPhrasesWithLen(ArrayList<int[]> sentences, int phraseSize){
		for (int i = 0; i < sentences.size(); i++) {
			int[] sentence = sentences.get(i);
			for (int index = 0; index < sentence.length; index++) {
				int phraseStartIndex = index;
				int phraseEndIndex = phraseStartIndex + phraseSize;
				if (phraseEndIndex > sentence.length)
					break;
				int[] phrase = new int[phraseSize];
				for (int s = 0; s < phraseSize; s++) {
					phrase[s] = sentence[index + s];
				}
				insertPhrase(phrase);
			}
		}
	}
	
	
//	Gets all phrases of a given lenght
	public void addPhrasesWithLenPrefixTrehs(ArrayList<int[]> sentences, int phraseSize, int prefixTresh){
		for (int i = 0; i < sentences.size(); i++) {
			int[] sentence = sentences.get(i);
			for (int index = 0; index < sentence.length; index++) {
				int phraseStartIndex = index;
				int phraseEndIndex = phraseStartIndex + phraseSize;
				if (phraseEndIndex > sentence.length)
					break;
				int[] phrase = new int[phraseSize];
				for (int s = 0; s < phraseSize; s++) {
					phrase[s] = sentence[index + s];
				}
				PhraseNode prev = getPhrase(phrase,phraseSize-1);
				if(prev != null && prev.count > prefixTresh){
					insertPhrase(phrase);
				}
			}
		}
	}
	
	//Add phrases of incremental lenght.
	//Remove phrases with count smaller than X
	//Only add phrases if count of smaller ones exist
	public void addPhrasesByLen(ArrayList<int[]> sentences, int prefixTresh){
			addPhrasesWithLen(sentences,1);
			for(int i = 2; i <= _maxPhraseLen; i++){
				addPhrasesWithLenPrefixTrehs(sentences, i, prefixTresh);
			}
	}
	
	public void writeToFile(PrintStream out){
		LinkedList<PhraseNode> stack = new LinkedList<PhraseNode>();
		TIntObjectHashMap<PhraseNode> children = _root.getChilds();
		for(int childrenKey : children.keys()){
			stack.addLast(children.get(childrenKey));
		}
		while(!stack.isEmpty()){
			PhraseNode node = stack.removeFirst();
			out.println(node._phraseId + " ||| " +  node.getCount() + " ||| " + intArrayToString(getPhraseIndexesById(node._phraseId)));
			TIntObjectHashMap<PhraseNode> nodeChildren = node.getChilds();
				for(int childrenKey : nodeChildren.keys()){
					stack.addLast(nodeChildren.get(childrenKey));
				}
		}
	}
	
	public void loadFromFile(BufferedReader reader) throws IOException{
		while(reader.ready()){
			String line = reader.readLine();
			String[] parts = line.split("\\s+\\|\\|\\|\\s+");
			int id = Integer.valueOf(parts[0]);
			int count = Integer.valueOf(parts[1]);
			String phrase = parts[2];
			insertPhrase(stringToIntArray(phrase), count, id);
		}
	}
	
	//loads from a .lm file format
	public void loadFromLMFile(BufferedReader reader, LanguageVocab vocab) throws IOException{
		while(reader.ready()){
			String line = reader.readLine();
			String[] words = line.split("\\t");
			if(words.length > 1){
				String ngram = words[1];
				double logprobability = Double.parseDouble(words[0]);
				String[] ngram_words = ngram.split("\\s+");
				int[] ngram_words_ints = vocab.getSentenceIds(ngram_words);
				PhraseNode node = insertPhraseRecursive(ngram_words_ints);
				node.setProbability(logprobability);
			}
		}
	}
	
	public void loadFromLMFileReverse(BufferedReader reader, LanguageVocab vocab) throws IOException{
		while(reader.ready()){
			String line = reader.readLine();
			String[] words = line.split("\\t");
			if(words.length > 1){
				String ngram = words[1];
				double logprobability = Double.parseDouble(words[0]);
				String[] ngram_words = ngram.split("\\s+");
				int[] ngram_words_ints = vocab.getSentenceIdsReverse(ngram_words);
				PhraseNode node = insertPhraseRecursive(ngram_words_ints);
				node.setProbability(logprobability);
			}
		}
	}
	
	public double getNGramProbability(int[] ngram){
		PhraseNode node = _root;
		int depth = ngram.length;
		for(int i = 0; i < depth; i++){
			node = node.getChild(ngram[i]);
			if(node == null){
				return node.getProbability();
			}
		}
		return 0;
	}

	
	public static void main(String[] args) throws IOException {
		String corpusDescription = args[0];
		int size = Integer.parseInt(args[1]); // 100k
		int maxSentenceSize = Integer.parseInt(args[2]); // 40
		int maxPhraseSize = Integer.parseInt(args[3]); // 40

		long initTime, initMem;
		long endMem;
		long endTime;
		initTime = System.currentTimeMillis();
		System.gc();
		initMem = (Runtime.getRuntime().totalMemory() - Runtime.getRuntime()
				.freeMemory())
				/ (1024 * 1024);
		BilingualCorpus corpus = BilingualCorpus.getCorpusFromFileDescription(
				corpusDescription, size, maxSentenceSize);

		endTime = System.currentTimeMillis();
		System.gc();
		System.gc();
		System.gc();
		endMem = (Runtime.getRuntime().totalMemory() - Runtime.getRuntime()
				.freeMemory())
				/ (1024 * 1024);
		System.out.println("Spent in Bulding corpus " + (endTime - initTime)
				/ 1000 + " mem " + (endMem - initMem));

		initTime = System.currentTimeMillis();
		System.gc();
		initMem = (Runtime.getRuntime().totalMemory() - Runtime.getRuntime()
				.freeMemory())
				/ (1024 * 1024);

		WordTrie sourcePhrasesVocab = new WordTrie(maxPhraseSize);
		sourcePhrasesVocab.addPhrasesAllSentences(corpus._trainSourceSentences);
		endTime = System.currentTimeMillis();
		System.gc();
		System.gc();
		System.gc();
		endMem = (Runtime.getRuntime().totalMemory() - Runtime.getRuntime()
				.freeMemory())
				/ (1024 * 1024);
		System.out.println("Spent in Bulding phrase Vocab "
				+ (endTime - initTime) / 1000 + " mem " + (endMem - initMem));

		sourcePhrasesVocab.compactTrie();
		endTime = System.currentTimeMillis();
		System.gc();
		System.gc();
		System.gc();
		endMem = (Runtime.getRuntime().totalMemory() - Runtime.getRuntime()
				.freeMemory())
				/ (1024 * 1024);
		System.out.println("Spent in Bulding phrase Vocab after compact"
				+ (endTime - initTime) / 1000 + " mem " + (endMem - initMem));
		System.out
				.println(("Different phraes" + sourcePhrasesVocab._phraseCounter));
		for(int i = 0; i < sourcePhrasesVocab._maxPhraseLen; i++){
			System.out.println((i+1) + " - " + sourcePhrasesVocab._phrasesByLen[i]);
		}
		
		System.out.println("Adding with tresh");
		initTime = System.currentTimeMillis();
		System.gc();
		initMem = (Runtime.getRuntime().totalMemory() - Runtime.getRuntime()
				.freeMemory())
				/ (1024 * 1024);

		WordTrie sourcePhrasesVocab2 = new WordTrie(maxPhraseSize);
		sourcePhrasesVocab2.addPhrasesByLen(corpus._trainSourceSentences,1);
		endTime = System.currentTimeMillis();
		System.gc();
		System.gc();
		System.gc();
		endMem = (Runtime.getRuntime().totalMemory() - Runtime.getRuntime()
				.freeMemory())
				/ (1024 * 1024);
		System.out.println("Spent in Bulding phrase Vocab with trehs"
				+ (endTime - initTime) / 1000 + " mem " + (endMem - initMem));
		
		sourcePhrasesVocab2.compactTrie();
		endTime = System.currentTimeMillis();
		System.gc();
		System.gc();
		System.gc();
		endMem = (Runtime.getRuntime().totalMemory() - Runtime.getRuntime()
				.freeMemory())
				/ (1024 * 1024);
		System.out.println("Spent in Bulding phrase Vocab with trehs after compacting"
				+ (endTime - initTime) / 1000 + " mem " + (endMem - initMem));
		System.out
				.println(("Different phraes" + sourcePhrasesVocab2._phraseCounter));
		for(int i = 0; i < sourcePhrasesVocab2._maxPhraseLen; i++){
			System.out.println((i+1) + " - " + sourcePhrasesVocab2._phrasesByLen[i]);
		}
				
	}

	public int getPhraseId(String string) {
		return getPhraseId(stringToIntArray(string));
	}

	public static int[] stringToIntArray(String string){
		String[] split = string.split("\\s+");
		int[] intConv = new int[split.length];
		for(int i = 0; i < intConv.length; i++){
			intConv[i] = Integer.parseInt(split[i]);
		}
		return intConv;
	}
	
	public static String intArrayToString(int[] intArray){
		String res = "";
		for(int i : intArray){
			res += i + " ";
		}
		return res;
	}

	public TIntObjectHashMap<PhraseNode> get_phrases() {
		return _phrases;
	}

	public PhraseNode get_root() {
		return _root;
	}	
}
